from pyabb import XYZPlatform
from pyati import ATISensor
from gelslimpi import GelSlimInterface, InputSource
from gelslimpi.visualize import qtcv

import cv2
import sys
from PyQt5 import QtWidgets
from scipy.interpolate import LinearNDInterpolator
from scipy.interpolate import RegularGridInterpolator as RGI

from datetime import datetime
from log_writer import LogWriter

from time import sleep
from threading import Thread
import numpy as np
from pathlib import Path
PROJ_DIR = Path(__file__).resolve().parent
CONF_DIR = PROJ_DIR / "configs"

from indenter import get_indenters

class DataCollector():
    """移动到任意像素坐标
    采集图片
    采集传感器数据
    下压"""
    def __init__(self, z_cont, cont_th1, cont_th2, log_dir=PROJ_DIR / f"runs/{datetime.now().strftime('%y_%m_%d_%H_%M_%S')}"):
        self.xyz = XYZPlatform("/dev/ttyUSB0", [-57, -119, -120])
        self.ati = ATISensor()
        self.ati.tare()
        self.source = InputSource("/raspicam_node32/image/compressed", "gelslim")
        self.gelslim = GelSlimInterface(configs=CONF_DIR/"configs3.yaml", input_source=self.source)
        self.init_interp()


        self.log_dir = Path(log_dir)
        self.log_dir.mkdir(exist_ok=True, parents=True)
        self.log_writer = LogWriter(self.log_dir/"log.csv",
                                    fields=["cls_name", "id", "x_pix", "y_pix", "depth_mm",
                                            "fx", "fy", "fz", "tx", "ty", "tz", "x_mm", "y_mm", "z_mm"])
        self.z_cont = z_cont#-38
        self.cont_th1 = cont_th1#-0.015
        self.cont_th2 = cont_th2#-0.035

    def show_img(self):
        while True:
            img = self.get_frame()
            qtcv.imshow("gelslim", img)
            key = qtcv.waitKey(200)
            if key == qtcv.Key.Key_Escape:
                break

    def get_frame(self):
        return self.gelslim.get_frame()[0]

    def get_ati(self):
        return self.ati.data

    def get_xyz(self):
        return self.xyz.x, self.xyz.y, self.xyz.z

    def init_interp(self, points_pix=None, points_xyz=None):
        """将像素坐标映射为实际坐标"""
        if points_pix is None:
            points_pix = np.array(
        [[45, 39], [315, 41], [591, 38], [580, 230], [577, 421], [313, 431], [43, 421], [41, 228], [314, 227], [549, 394]]
                                  )
        else:
            points_pix = np.array(points_pix)
        if points_xyz is None:
            points_xyz = np.array(
[[9.625, -63.286], [1.645, -63.56], [-6.481, -63.563], [-6.401, -57.836], [-6.51, -52.211], [1.53, -52.025], [9.417, -51.936], [9.67, -57.657], [1.631, -57.898], [-5.61, -53.083]]
)
        else:
            points_xyz = np.array(points_xyz)
        # 用非结构化插值获得网格数据
        interp = LinearNDInterpolator(points=points_pix, values=points_xyz)
        ymin = np.max(points_pix[0:3,1])
        xmax = np.min(points_pix[2:5,0])
        ymax = np.min(points_pix[4:7,1])
        xmin = np.max([points_pix[6,0], points_pix[7,0], points_pix[0,0]])
        x_span = np.linspace(xmin, xmax, 9)
        y_span = np.linspace(ymin, ymax, 9)
        X, Y = np.meshgrid(x_span, y_span, indexing="ij")
        grids = interp(X, Y)
        self.interp = RGI((x_span, y_span), grids, bounds_error=False, fill_value=None)

    def move(self, pix_x, pix_y, bias=(0,0)):
        # 移动到接触高度上方， 力传感器清零
        self.xyz.move_abs(z = self.z_cont+2, vel=40)  # z 提升到 -26
        self.ati.tare()
        # xy 移动
        x, y = self.interp((pix_x, pix_y)) + bias
        print(f"move to [{pix_x:.2f}, {pix_y:.2f}]pix, [{x:.2f}, {y:.2f}]mm")
        self.xyz.move_abs(x=x, y=y, vel=40)
        # 向下移动
        self.xyz.move_rel(z=-1.3)  # for speed
        self.down_util_contact(pix_x, pix_y)


    def log_data(self, cls_name, id, px, py, z_mm):
        """记录数据"""
        img, stamp, raw_img = self.gelslim.get_frame_all()
        img_path = str(self.log_dir / f"{cls_name}_{id:>04d}.jpg")
        raw_img_path = str(self.log_dir / f"raw_{cls_name}_{id:>04d}.jpg")
        data = self.ati.data
        self.log_writer.log(cls_name, id, px, py, z_mm, *list(data),
                            self.xyz.x, self.xyz.y, self.xyz.z, echo=True)
        cv2.imwrite(img_path, img)
        cv2.imwrite(raw_img_path, raw_img)
        qtcv.imshow("contact", img)
        qtcv.waitKey(5)

    # def down_util_contact(self):
    #     is_contact = False
    #     while not is_contact:
    #         # move down slowly
    #         self.xyz.move_rel(z=-0.01, vel=50)
    #         sleep(0.05)
    #         img = self.get_frame()
    #         qtcv.imshow("contact", img)
    #         key = qtcv.waitKey(5)
    #         if key == qtcv.Key.Key_Escape:
    #             exit()

    #         # 检测法向力
    #         fz1 = self.get_ati()[2]
    #         if fz1 <= self.cont_th1:
    #             # 再向下 0.1 确认是否接触
    #             self.xyz.move_rel(z=-0.1, vel=30)
    #             fz2 = self.get_ati()[2]
    #             if fz2 <= self.cont_th2:
    #                 self.z_cont = self.xyz.z+0.1   # 更新上次接触的高度
    #                 print(f"contact detected. z_cont: {self.z_cont}mm | fz1: {fz1} | fz2: {fz2}")
    #                 is_contact = True
    #             self.xyz.move_rel(z=0.1, vel=30)


    def down_util_contact(self, pix_x, pix_y):
        cont_cnt = 0  # 检测到接触的次数
        step = [0.15, 0.01]
        while cont_cnt < len(step):
            # move down slowly
            self.xyz.move_rel(z=-step[cont_cnt], vel=40)
            sleep(0.05)
            img = self.get_frame()
            cv2.circle(img, (pix_x, pix_y), 5, (0,0,255), 1)
            qtcv.imshow("contact", img)
            key = qtcv.waitKey(5)
            if key == qtcv.Key.Key_Escape:
                exit()

            # 检测法向力
            fz1 = self.get_ati()[2]
            if fz1 <= self.cont_th1:
                # 再向下 0.1 确认是否接触
                self.xyz.move_rel(z=-0.1, vel=40)
                fz2 = self.get_ati()[2]
                is_contact = fz2 <= self.cont_th2
                self.xyz.move_rel(z=0.1, vel=40)
                if is_contact and cont_cnt == len(step)-1:  # 最后一次收敛
                    self.z_cont = self.xyz.z
                    print(f"contact detected. z_cont: {self.z_cont}mm | fz1: {fz1} | fz2: {fz2}")
                    break
                elif is_contact:  # 不是最后一次接触
                    self.xyz.move_rel(z=step[cont_cnt]+0.02, vel=40)
                    cont_cnt += 1

if __name__ == '__main__':
    indent_names, indent_bias = get_indenters("flat")  # 压头类型
    points_num = 100  # 点数
    depth_num = 3 # 每个点采样的不同深度数量
    depth_range = (0.2, 1.1)  # 深度范围
    xy_pad = 20  # 采样点范围不包括边缘, unit: pix

    dc = DataCollector(z_cont=-239, cont_th1=-0.015, cont_th2=-0.035)

    for idx, indent_name in enumerate(indent_names):
        # if idx < 6 :
        #     continue
        bias = indent_bias[idx]
        points_cnt = 0
        id = 0
        # if idx==6:
        #     points_cnt = 92
        #     id = 276

        while points_cnt < points_num:
            x = np.random.randint(xy_pad, 640-xy_pad)  # 随机采样一个点
            y = np.random.randint(xy_pad, 480-xy_pad)
            depths = np.random.uniform(depth_range[0], depth_range[1], depth_num)

            ret = dc.move(x, y, bias)
            points_cnt += 1

            z = dc.z_cont  # 接触深度
            for depth in depths:
                dc.xyz.move_abs(z=z-depth)
                sleep(0.8)
                dc.log_data(indent_name, id, x, y, depth)
                id += 1
        dc.xyz.move_abs(z=z+10)
        # 更换压头，不同压头高度不一， 为了不发生碰撞， 将碰撞高度上调
        dc.z_cont += 1.5
